import os
from PIL.Image import blend
import cupy as cp
import numpy as np
import matplotlib.pyplot as plt
import cv2
import concurrent.futures
import time
import glob

python_file = os.path.dirname(__file__)
results = os.path.join(python_file, 'results')
masks_path = os.path.join(results, 'mask-original')
original = os.path.join(python_file, 'original')
results_path = os.path.join(results, 'blend-fitted')

def make_mask(mask):
    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    ret, mask = cv2.threshold(mask, 0, 1, cv2.THRESH_BINARY)
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    area = cp.array([cv2.contourArea(_) for _ in contours])
    max_area_id = cp.argmax(area)
    for _ in range(len(contours)):
        if _ != max_area_id:
            cv2.fillPoly(mask, [contours[_]], 0)
    return mask

def get_var(mask_):
    try:
        mask = cp.array(mask_)
        sumx = cp.sum(mask, axis=0)
        indexes = cp.argwhere(sumx != 0)
        xr = cp.min(indexes)
        xl = cp.max(indexes)
        a = int((xr - xl) * 0.5)
        dx = int((xr + xl) * 0.5)
        axisy = mask[:, dx]
        indexes = cp.argwhere(axisy != 0)
        yu = cp.min(indexes)
        yd = cp.max(indexes)
        b = int((yd - yu) * 0.5)
        dy = int((yd + yu) * 0.5)
    except:
        mask = np.array(mask_)
        sumx = np.sum(mask, axis=0)
        indexes = np.argwhere(sumx != 0)
        xr = np.min(indexes)
        xl = np.max(indexes)
        a = int((xr - xl) * 0.5)
        dx = int((xr + xl) * 0.5)
        axisy = mask[:, dx]
        indexes = np.argwhere(axisy != 0)
        yu = np.min(indexes)
        yd = np.max(indexes)
        b = int((yd - yu) * 0.5)
        dy = int((yd + yu) * 0.5)
    return [a, b, dx, dy]


def iou(mask_, variables):
    union = 0.
    inter = 0.
    try:
        [a, b, dx, dy, k] = cp.array(variables)
        mask = cp.array(mask_)
        [ymax, xmax] = cp.shape(mask)
        
        x = cp.arange(xmax)
        y = cp.arange(ymax)
        x, y = cp.meshgrid(x, y)

        egg = cp.square((x - dx)/a) + cp.square((y - dy)/b) * cp.exp(k * (x - dx))
        egg = cp.where(egg <= 1, 1, 0)
        inter = egg * mask
        inter_ = cp.sum(inter)
        union = egg + mask
        union = cp.where(union != 0, 1, 0)
        union_ = cp.sum(union)
    except:
        [a, b, dx, dy, k] = variables
        mask = np.array(mask_)
        [ymax, xmax] = np.shape(mask)
        
        x = np.arange(xmax)
        y = np.arange(ymax)
        x, y = np.meshgrid(x, y)

        egg = np.square((x - dx)/a) + np.square((y - dy)/b) * np.exp(k * (x - dx))
        egg = np.where(egg <= 1, 1, 0)
        inter = egg * mask
        inter_ = np.sum(inter)
        union = egg + mask
        union = np.where(union != 0, 1, 0)
        union_ = np.sum(union)

    '''
    cv2.imshow('result', np.hstack((np.vstack((egg, mask)), np.vstack((inter, union)))))
    cv2.waitKey()
    cv2.destroyAllWindows()
    '''
    return float(inter_ / union_)

def plotegg(variables):
    [a, b, dx, dy, k] = variables
    X = np.linspace(-a, a, 1000)
    Y = np.sqrt((1 - np.square(X/a)) / np.exp(k * X)) * b
    plt.plot(X+dx, Y+dy, color = 'yellow')
    plt.plot(X+dx, -Y+dy, color = 'yellow')
    plt.axis('off')
    X = np.linspace(-a, a, 20)
    Y = np.sqrt((1 - np.square(X/a)) / np.exp(k * X)) * b
    X_ = (X + dx)[::-1].tolist()
    X = (X + dx).tolist()
    X += X_
    Y_ = (-Y)[::-1].tolist()
    Y = (Y).tolist()
    Y += Y_
    return np.array([X, Y]).T

def get_egg(mask, pic, save_path, pic_path):
    mask = np.array(make_mask(mask), dtype=float)

    variables = get_var(mask)
    variables.append(0.2/30)
    variables = np.array(variables)
    learn_rate = 0.3
    grad = 0.1/600
    iou_0 = iou(mask, variables)
    iou_p = iou_0
    iou_his = [0]
    k_his = [0]

    while True:
        
        variables[4] += learn_rate*grad
        iou_p = iou(mask, variables)
        
        if iou_p > iou_0:
            grad = iou_p - iou_0
        else:
            variables[4] -= 1.25*learn_rate*grad/iou_0
            iou_p = iou(mask, variables)
            grad = iou_p - iou_0
        iou_0 = iou_p
        iou_his.append(iou_0)
        k_his.append(variables[4])
        if iou_his[-1] <= iou_his[-2]:
            variables[4] = k_his[-2]
            break
        if iou_0 > 0.9:
            break
    if variables[4] <=0:
        print('\nError k in ' + pic_path)
        write_type = "a" if os.path.exists("error pic list.txt") else "w"
        with open("error pic list.txt", write_type) as f:
            f.write(pic_path + "\n") 
    plt.cla()
    plt.imshow(pic)
    np.savetxt(save_path.replace('jpg', 'txt'), plotegg(variables), fmt='%.15f')
    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)

def main(pic):
    
    mask_path = os.path.join(masks_path, pic)
    pic_path = os.path.join(original, pic)
    result_path = os.path.join(results_path, pic)
    
    try:
        mask = cv2.imread(mask_path)
        pic =  np.array(cv2.cvtColor(cv2.imread(pic_path), cv2.COLOR_BGR2RGB))
        get_egg(mask, pic, result_path, pic_path)
    except Exception as e:
        print('\nERROR:%s'%pic)
        print(repr(e))
    
def make_dirs(dirs):
    for dir in dirs:
        if not os.path.exists(dir):
            os.makedirs(dir)

if __name__ == '__main__':
    make_dirs([results, masks_path, results_path])
    start = time.time()
    pics = glob.glob(os.path.join(original, '*.jpg'))
    with concurrent.futures.ProcessPoolExecutor() as executor:
        executor.map(main, pics)
    end = time.time()
    print('TIME: ', end - start)